#include "../Base.h"

NTSTATUS GetModuleX86(PEPROCESS process,PPEB32 peb32,PUNICODE_STRING moduleName, PULONG_PTR imageBase, SIZE_T* imageSize) {

	SIZE_T retSize = 0;
	NTSTATUS status;
	ULONG_PTR module = 0;
	status = MmCopyVirtualMemory(process, peb32, process, peb32, 0x1, UserMode, &retSize);
	if (!NT_SUCCESS(status)) {
		return STATUS_UNSUCCESSFUL;
	}
	PPEB_LDR_DATA32 pebLdr =  (PPEB_LDR_DATA32)(peb32->Ldr);
	LIST_ENTRY32 currentEntry = pebLdr->InLoadOrderModuleList;
	PLDR_DATA_TABLE_ENTRY32 nextEntry = (PLDR_DATA_TABLE_ENTRY32)(currentEntry.Flink);
	while (nextEntry != (PVOID)&currentEntry) {


		UNICODE_STRING baseDllName = nextEntry->BaseDllName;
		if (RtlCompareUnicodeString(&baseDllName, moduleName, TRUE) == 0) {
			if(imageBase)
			*(ULONG_PTR*)imageBase = nextEntry->DllBase;
			if(imageSize)
			*imageSize = nextEntry->SizeOfImage;
			LOG("[Image Info] Image base = %llx  Image Size = %llx \r\n", nextEntry->DllBase, nextEntry->SizeOfImage);
			return STATUS_SUCCESS;
		}

		nextEntry = (PLDR_DATA_TABLE_ENTRY32)nextEntry->InLoadOrderLinks.Flink;
	}
	return STATUS_UNSUCCESSFUL;
}
NTSTATUS GetModuleX64(PEPROCESS process, PPEB peb, PUNICODE_STRING moduleName, PULONG_PTR imageBase, SIZE_T* imageSize) {

	SIZE_T retSize = 0;
	NTSTATUS status;
	ULONG_PTR module = 0;
	status = MmCopyVirtualMemory(process, peb, process, peb, 8, UserMode, &retSize);
	PPEB_LDR_DATA pebLdr = (PPEB_LDR_DATA)(peb->Ldr);
	LIST_ENTRY currentEntry = pebLdr->InLoadOrderModuleList;
	PLDR_DATA_TABLE_ENTRY nextEntry = (PLDR_DATA_TABLE_ENTRY)(currentEntry.Flink);
	while (nextEntry != (PVOID)&currentEntry) {


		PWCH baseDllName = nextEntry->BaseDllName.Buffer;
		UNICODE_STRING uBaseName;
		RtlInitUnicodeString(&uBaseName, baseDllName);
		
		if (RtlCompareUnicodeString(&uBaseName, moduleName, TRUE) == 0) {
			//LOG("baseDllName = %ls", uBaseName.Buffer);
			if (imageBase)
				*(PULONG_PTR)imageBase = (ULONG_PTR)nextEntry->DllBase;
			if (imageSize)
				*imageSize = nextEntry->SizeOfImage;
			//LOG("[Image Info] Image base = %llx  Image Size = %llx", nextEntry->DllBase, nextEntry->SizeOfImage);
			return STATUS_SUCCESS;
		}

		nextEntry = (PLDR_DATA_TABLE_ENTRY)nextEntry->InLoadOrderLinks.Flink;
	}
	return STATUS_UNSUCCESSFUL;
}
NTSTATUS GetModuleInfo(HANDLE pid, char* moduleName, PULONG_PTR imageBase, SIZE_T* imageSize)
{
	PEPROCESS process;
	KAPC_STATE apc_state = { 0 };
	if (pid < 0 || strlen(moduleName) == 0) {
		return STATUS_UNSUCCESSFUL;
	}
	NTSTATUS status = PsLookupProcessByProcessId((HANDLE)pid, (PEPROCESS*)&process);
	if (!NT_SUCCESS(status)) {
		LOG("���ҽ��̴���\r\n");
		return STATUS_UNSUCCESSFUL;
	}
	ANSI_STRING aMouduleName = { 0 };
	UNICODE_STRING uModuleName = { 0 };
	RtlInitAnsiString(&aMouduleName, moduleName);
	status = RtlAnsiStringToUnicodeString(&uModuleName, &aMouduleName, TRUE);
	if (!NT_SUCCESS(status)) {
		LOG("ת���ַ�������\r\n");
		return STATUS_UNSUCCESSFUL;
	}


	KeStackAttachProcess((PRKPROCESS)process, &apc_state);

	PPEB peb = PsGetProcessPeb(process);
	PPEB32 peb32  =  (PEB32*)PsGetProcessWow64Process((PEPROCESS)process);
	if (peb32) {
		status = GetModuleX86((PEPROCESS)process, peb32,&uModuleName,imageBase,imageSize);
	}
	else {
		status = GetModuleX64((PEPROCESS)process, peb, &uModuleName, imageBase, imageSize);
	}
	if (!NT_SUCCESS(status)) {
		return STATUS_UNSUCCESSFUL;
	}


	KeUnstackDetachProcess(&apc_state);
	return STATUS_SUCCESS;
}








//系统模块信息

ULONG_PTR QuerySystemModuleInfo(char* moduleName, PULONG_PTR moduleSize) {

	ULONG_PTR;
	RTL_PROCESS_MODULES rtlModules = { 0 };
	PRTL_PROCESS_MODULES SystemModules = &rtlModules;
	ULONG retLen = 0;
	BOOLEAN isAllocate = FALSE;
	NTSTATUS status = ZwQuerySystemInformation(
		SystemModuleInformation,
		SystemModules, sizeof(RTL_PROCESS_MODULES),
		&retLen
	);
	if (status == STATUS_INFO_LENGTH_MISMATCH) {
		SystemModules = (PRTL_PROCESS_MODULES)ExAllocatePoolWithTag(PagedPool, retLen + sizeof(RTL_PROCESS_MODULES), 'yang');
		memset(SystemModules, 0, retLen + sizeof(RTL_PROCESS_MODULES));
		status = ZwQuerySystemInformation(
			SystemModuleInformation,
			SystemModules, retLen + sizeof(RTL_PROCESS_MODULES),
			&retLen);
		if (!NT_SUCCESS(status)) {
			ExFreePoolWithTag(SystemModules, retLen + sizeof(RTL_PROCESS_MODULES));
			return 0;
		}
		isAllocate = TRUE;
	}


	ULONG_PTR moduleBase = 0;
	do {

	

		if (_stricmp(moduleName, "ntkrnlpa.exe") == 0 || _stricmp(moduleName, "ntoskrnl.exe") == 0) {

			PRTL_PROCESS_MODULE_INFORMATION moduleInfo = &SystemModules->Modules[0];
			moduleBase = (ULONG_PTR)moduleInfo->ImageBase;
			if (moduleSize) *moduleSize = moduleInfo->ImageSize;
			break;
		}


		char* kernelModuleName = (char*)ExAllocatePoolWithTag(PagedPool, strlen((char*)moduleName) + 1, 'yang');
		memset(kernelModuleName, 0, strlen((char*)moduleName) + 1);
		memcpy(kernelModuleName, moduleName, strlen(moduleName) + 1);
		_strupr(kernelModuleName);

		for (int i = 1; i < SystemModules->NumberOfModules; i++) {
			PRTL_PROCESS_MODULE_INFORMATION moduleInfo = &SystemModules->Modules[i];

			//LOG("[MODULE]:baseName = %s \r\n", moduleInfo->FullPathName + moduleInfo->OffsetToFileName);
			char* pathName = _strupr((char*)moduleInfo->FullPathName);

			if (strstr(pathName, kernelModuleName)) {
				moduleBase = (ULONG_PTR)moduleInfo->ImageBase;
				if (moduleSize) *moduleSize = moduleInfo->ImageSize;
				break;
			}
		}
		if (kernelModuleName) {
			ExFreePoolWithTag(kernelModuleName, 'yang');
		}
	} while (0);

	if (isAllocate) {
		ExFreePoolWithTag(SystemModules, 'yang');
	}

	return moduleBase;
}